"""
文本嵌入向量相似度 测试
1.嵌入模型
2.余弦相似度公式
3.文本向量数据数据 列表
4.查询和被查询的数据向量化
5.向量相似度比较
"""
from numpy import dot
from numpy.linalg import norm
from model_utils import getEmd

# 1.嵌入模型
embedding_model = getEmd()

# 2.余弦相似度公式
def cos_sim(a, b):
    """
    余弦相似度 -- 越大越相似
    :param a:
    :param b:
    :return:
    """
    return dot(a, b) / (norm(a) * norm(b))

# 3.文本向量数据
def get_vectors_data(texts):
    list_vec = list()
    for t in texts:
        embedding = embedding_model.embed_query(t)
        list_vec.append(embedding)
    return list_vec

# 4.向量数据列表
datas = [
    "基本信息",
    "姓名：马云（英文名：Jack Ma）",
    "出生日期：1964年9月10日",
    "出生地：中国浙江省杭州市",
    "业院校：杭州师范大学（原杭州师范学院），英语专业",
    "花名：风清扬（源自金庸小说《笑傲江湖》中的神秘高手风清扬）"
]

query = "马云的外号是什么"

# 4.查询和被查询的数据向量化
query_vec = get_vectors_data([query])[0]
doc_vecs = get_vectors_data(datas)

# 5.向量比较
for vec in doc_vecs:
    print(cos_sim(query_vec, vec))